package com.lagou.cloud.gateway.filter;

import com.lagou.cloud.gateway.dao.RequestDAO;
import com.lagou.cloud.gateway.pojo.Request;
import com.lagou.common.util.CommonUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.sql.Timestamp;

@Slf4j
@Component
public class RequestLimitFilter implements GlobalFilter, Ordered {

    @Value("${custom.limit.duration}")
    private Long limitDuration;

    @Value("${custom.limit.times}")
    private Long limitTimes;

    @Autowired
    private RequestDAO requestDAO;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

        ServerHttpRequest serverHttpRequest = exchange.getRequest();
        ServerHttpResponse serverHttpResponse = exchange.getResponse();

        String clientIp = serverHttpRequest.getRemoteAddress().getHostString();
        String path = serverHttpRequest.getPath().toString();
        Timestamp now = CommonUtils.now();
        Timestamp limitDurationTime = CommonUtils.expiresTime(now, limitDuration * -1);

        Long requestCount = requestDAO.countByIpAndPathAndCreateTimeGreaterThanEqual(clientIp, path, limitDurationTime);
        if (requestCount > limitTimes) {
            serverHttpResponse.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
            log.warn("{}请求超过限制", clientIp);

            String message = "请求超过限制，请稍后再试";
            DataBuffer buffer = serverHttpResponse.bufferFactory().wrap(message.getBytes());

            return serverHttpResponse.writeWith(Mono.just(buffer));
        }

        Request request = new Request();
        request.setIp(clientIp);
        request.setPath(path);
        request.setCreateTime(now);
        requestDAO.save(request);

        return chain.filter(exchange);
    }

    @Override
    public int getOrder() {
        return 1;
    }
}
